import visdom


class GANVisualization:
    def __init__(self, env:str):
        self.viz = visdom.Visdom(env=env)
        self.loss_g = []
        self.loss_d = []

    def vis_samples(self, g, d, noise_func, caption:str='dcgan'):
        tensor_gen = g(noise_func())
        n, c, h, w = tensor_gen.shape
        self.viz.images(tensor_gen, opts={
            'nrow': 6,
            'caption': caption,
            'win': 'samples'
        })

    def vis_loss(self, loss_g, loss_d):
        self.loss_g.append(loss_g)
        self.loss_d.append(loss_d)
        self.viz.line(X=list(range(len(self.loss_g))), Y=self.loss_g, win='loss g', name='loss g',
            opts={
                'title': 'loss g',
            }
        )
        self.viz.line(X=list(range(len(self.loss_d))), Y=self.loss_d, win='loss d', name='loss d', 
            opts={
                'title': 'loss d',
            }
        )
